from sklearn.datasets import load_iris
from flower import *
from flowerinput import *
import knn
import random


def test():
    data = load_iris()

    setosa = Flower(0, 'Setosa')
    versicolour = Flower(1, 'Versicolour')
    virginica = Flower(2, 'Virginica')

    totalLen = len(data.data)
    testLen = 30

    dataSet = []
    for i in range(totalLen):
        dataSet.append([data.data[i], data.target[i]])
    random.shuffle(dataSet)

    trainSet = []
    testSet = []
    for i in range(totalLen):
        input = FlowerInput(dataSet[i][0])
        target = None
        if dataSet[i][1] == 0:
            target = setosa
        elif dataSet[i][1] == 1:
            target = versicolour
        elif dataSet[i][1] == 2:
            target = virginica
        if i < totalLen - testLen:
            trainSet.append([input, target])
        else:
            testSet.append([input, target])
    
    k = 7
    knn.test(trainSet, testSet, k)

test()
